import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
from model import MyModel
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets,losses,layers,metrics,Sequential,optimizers

def preprocess(x, y):
    """
    x is a simple image, not a batch
    """
    x = tf.cast(x, dtype=tf.float32) / 255.
    x = tf.reshape(x, [28*28])
    y = tf.cast(y, dtype=tf.int32)
    y = tf.one_hot(y, depth=10)
    return x,y

batchsz = 128
(x, y), (x_test, y_test) = datasets.mnist.load_data()#加载图片
print('datasets:', x.shape, y.shape, x.min(), x.max())


db = tf.data.Dataset.from_tensor_slices((x,y))
db = db.map(preprocess).shuffle(60000).batch(batchsz)
ds_val = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = ds_val.map(preprocess).batch(batchsz)

sample = next(iter(db))
print(sample[0].shape, sample[1].shape)


network = MyModel()

network.compile(
    optimizer=optimizers.Adam(lr=0.01),
    loss=tf.losses.CategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

network.fit(db,
    epochs=3,
    validation_data=db_test,
    validation_freq=2
)

network.evaluate(db_test)

network.save_weights('weights.ckpt')
print('saving weights.')
tf.saved_model.save(network,'savedmodel/')
print('saving model')

sample = next(iter(db_test))
x = sample[0]
y = sample[1]
pred = network.predict(x)

#convert back to number
y = tf.argmax(y, axis=1)
pred = tf.argmax(pred, axis=1)

print(pred)
print(y)

